import cv2
from utils import block_view, blkproc, IDCT2, PCAdenoising
import numpy as np
import matplotlib.pyplot as plt
from skimage import io, measure, util


def check_block():
    img_path = 'original/lena.png'
    img_arr = cv2.imread(img_path, 0)
    print img_arr.shape
    view, shape = block_view(img_arr, (8, 8))
    print np.asarray(view).shape
    a = blkproc(img_arr, (8, 8), IDCT2)
    print a.shape
    A = np.uint8(PCAdenoising(img_arr, (8, 8), 2))
    print measure.compare_psnr(im_true=img_arr, im_test=A, dynamic_range=255)
    fig = plt.figure()
    io.imshow(A)
    fig.savefig('test.png')


def add_noise():
    img_path = 'original/lena.png'
    img_arr = cv2.imread(img_path, 0)
    img_arr_noise = util.random_noise(img_arr, 'gaussian', var=0.05)
    cv2.imwrite('test/gaussian.png', np.asarray(img_arr_noise * 255., np.uint8))

    img_arr_salt = util.random_noise(img_arr, 's&p', amount=0.1)
    cv2.imwrite('test/slat_pepper.png', np.asarray(img_arr_salt * 255., np.uint8))

    img_arr_speckle = util.random_noise(img_arr, 'speckle', var=0.1)
    cv2.imwrite('test/speckle.png', np.asarray(img_arr_speckle * 255., np.uint8))
